package main

import (
	"context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
	"log"
	"net"
	"os"
	"path"
	"strings"
)

const (
	resourceName = "example.com/hardware-device"
	socketName   = "hardware.sock"
	kubeletSocket = "kubelet.sock"  // kubelet 的注册服务端点
)

type DemoDevicePlugin struct {
	devices []*pluginapi.Device
	stop    chan struct{}
}

func NewDemoDevicePlugin() *DemoDevicePlugin {
	return &DemoDevicePlugin{
		devices: []*pluginapi.Device{
			{ID: "Device-1", Health: pluginapi.Healthy},
			{ID: "Device-2", Health: pluginapi.Healthy},
		},
		stop: make(chan struct{}),
	}
}

func (m *DemoDevicePlugin) ListAndWatch(e *pluginapi.Empty, s pluginapi.DevicePlugin_ListAndWatchServer) error {
	resp := &pluginapi.ListAndWatchResponse{Devices: m.devices}
	if err := s.Send(resp); err != nil {
		return err
	}

	<-m.stop
	return nil
}

func (m *DemoDevicePlugin) Allocate(ctx context.Context, req *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
	responses := &pluginapi.AllocateResponse{}
	for _, containerReq := range req.ContainerRequests {
		response := &pluginapi.ContainerAllocateResponse{
			Envs: map[string]string{
				"HARDWARE_DEVICE_ID": containerReq.DevicesIDs[0],
			},
		}
		responses.ContainerResponses = append(responses.ContainerResponses, response)
	}
	return responses, nil
}

func (m *DemoDevicePlugin) GetDevicePluginOptions(context.Context, *pluginapi.Empty) (*pluginapi.DevicePluginOptions, error) {
	return &pluginapi.DevicePluginOptions{}, nil
}

func (m *DemoDevicePlugin) PreStartContainer(context.Context, *pluginapi.PreStartContainerRequest) (*pluginapi.PreStartContainerResponse, error) {
	return &pluginapi.PreStartContainerResponse{}, nil
}

func (m *DemoDevicePlugin) GetPreferredAllocation(ctx context.Context, req *pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) {
	// 简单实现：直接返回请求的设备
	return &pluginapi.PreferredAllocationResponse{
		ContainerResponses: []*pluginapi.ContainerPreferredAllocationResponse{
			{DeviceIDs: req.ContainerRequests[0].AvailableDeviceIDs},
		},
	}, nil
}

func main() {
	plugin := NewDemoDevicePlugin()
	server := grpc.NewServer()
	pluginapi.RegisterDevicePluginServer(server, plugin)

	socketPath := path.Join(pluginapi.DevicePluginPath, socketName)
	if err := os.Remove(socketPath); err != nil && !os.IsNotExist(err) {
		log.Fatal(err)
	}

	listener, err := net.Listen("unix", socketPath)
	if err != nil {
		log.Fatal(err)
	}

	go func() {
		if err := server.Serve(listener); err != nil {
			log.Fatal(err)
		}
	}()


	// 连接到 kubelet 的注册服务（注意使用不同的路径）
	kubeletSocketPath := path.Join(pluginapi.DevicePluginPath, kubeletSocket)
	conn, err := grpc.Dial("unix://"+kubeletSocketPath, 
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) {
			// 去除协议前缀
			addr = strings.TrimPrefix(addr, "unix://")
			return net.Dial("unix", addr)
		}))
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()

	client := pluginapi.NewRegistrationClient(conn)
	_, err = client.Register(context.Background(),
		&pluginapi.RegisterRequest{
			Version:      pluginapi.Version,
			Endpoint:     socketName,
			ResourceName: resourceName,
		},
	)
	if err != nil {
		log.Fatal(err)
	}

	log.Println("Device plugin registered")
	select {}
}
